tests/attention/test_rope.py (1)

1390-1590: Add enable_pdl coverage to this new padding regression test.

Lines 1392-1589 only exercise the default path. Please parameterize enable_pdl and pass it into the fused call so padding behavior is validated under the programmatic dependent launch mode too.

Proposed test update

 `@pytest.mark.parametrize`("kv_layout", ["NHD", "HND"])
 `@pytest.mark.parametrize`("page_size", [16])
+@pytest.mark.parametrize("enable_pdl", [True, False])
 def test_rope_quantize_fp8_append_paged_kv_cache_padding(
@@
     kv_layout,
     page_size,
+    enable_pdl,
 ):
@@
     flashinfer.rope.rope_quantize_fp8_append_paged_kv_cache(
@@
         quant_scale_kv=1.0,
         is_neox=False,
+        enable_pdl=enable_pdl,
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/attention/test_rope.py` around lines 1390 - 1590, The test
test_rope_quantize_fp8_append_paged_kv_cache currently only runs the default
path; add a pytest parameterization for enable_pdl (e.g.,
`@pytest.mark.parametrize`("enable_pdl",[False,True]) and add enable_pdl to the
test signature) and pass enable_pdl into the fused call
flashinfer.rope.rope_quantize_fp8_append_paged_kv_cache(enable_pdl=enable_pdl,
... ) so the padding behavior is validated under programmatic dependent launch
mode as well.

include/flashinfer/pos_enc.cuh (1)

862-868: Make padding-sentinel check type-safe across PagedKVIdType instantiations.

Line 863 uses batch_indices[idx] >= 0, which is only safe when PagedKVIdType is signed. If it is ever unsigned, the padding sentinel -1 becomes the maximum value and this branch incorrectly passes, leading to invalid indptr indexing.

🔧 Proposed fix

-    // skip padding tokens with batch_indices < 0
-    if (batch_indices[idx] >= 0) {
+    constexpr PagedKVIdType kPaddingSentinel = static_cast<PagedKVIdType>(-1);
+    const PagedKVIdType batch_idx = batch_indices[idx];
+    if (batch_idx != kPaddingSentinel) {
       // Compute page location for this token
       uint32_t page_iter, entry_idx;
       paged_kv_like.page_size.divmod(
-          paged_kv_like.indptr[batch_indices[idx]] * paged_kv_like.page_size + positions[idx],
+          paged_kv_like.indptr[batch_idx] * paged_kv_like.page_size + positions[idx],
           page_iter, entry_idx);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@include/flashinfer/pos_enc.cuh` around lines 862 - 868, The current padding
check uses "batch_indices[idx] >= 0" which breaks for unsigned PagedKVIdType;
change it to a type-safe sentinel comparison (e.g., compare against a named
padding sentinel value) so you explicitly test for the padding marker instead of
signedness. Replace the condition in the block using batch_indices, positions
and paged_kv_like (the if around batch_indices[idx] >= 0 that precedes
paged_kv_like.indptr[...] and page_size.divmod calls) with a check like
"batch_indices[idx] != static_cast<PagedKVIdType>(-1)" or, better, introduce a
constexpr PagedKVIdType kPaddingSentinel = static_cast<PagedKVIdType>(-1) and
use "batch_indices[idx] != kPaddingSentinel" to ensure correct behavior for both
signed and unsigned instantiations.

include/flashinfer/pos_enc.cuh (1)

862-863: Consider adding a compile-time assertion to document the signed type requirement for the padding check.

The code assumes PagedKVIdType is signed to make the check batch_indices[idx] >= 0 work correctly for detecting padding tokens (marked with -1). While the type is hardcoded to int32_t in the JIT binding at csrc/rope.cu (lines 574, 605) and in Python via torch.int(), a static assertion would document this requirement and guard against accidental misuse:

static_assert(std::is_signed<PagedKVIdType>::value,
              "PagedKVIdType must be signed for padding token detection (batch_indices < 0)");

This defensive check aligns with the kernel's assumptions and improves long-term maintainability without adding runtime overhead.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@include/flashinfer/pos_enc.cuh` around lines 862 - 863, Add a compile-time
assertion that PagedKVIdType is a signed type to document and enforce the
kernel's assumption used by the padding check `batch_indices[idx] >= 0`; insert
a static_assert using `std::is_signed<PagedKVIdType>::value` (e.g., near the
typedef/using of PagedKVIdType or at the top of the kernel in pos_enc.cuh before
the `batch_indices` usage) with a clear message like "PagedKVIdType must be
signed for padding token detection (batch_indices < 0)"; this is purely
compile-time and has no runtime overhead but prevents accidental unsigned types
from being used.